from physics.vector.array import Vector4Array as V4Array
from physics.const import KCHARGED_MASS_GEV as KMASS
from physics.const import PICHARGED_MASS_GEV as PIMASS

import numpy

from charmicat import ColumnCalculator

from cmix.parameters import paramdef

def four_vectors(r):
    v = {}
    v['kp'] = V4Array(r['kp_px'],
                      r['kp_py'],
                      r['kp_pz'],
                      m=KMASS)
    v['km'] = V4Array(r['km_px'],
                      r['km_py'],
                      r['km_pz'],
                      m=KMASS)

    v['pip'] = V4Array(r['pip_px'],
                       r['pip_py'],
                       r['pip_pz'],
                       m=PIMASS)
    v['pim'] = V4Array(r['pim_px'],
                       r['pim_py'],
                       r['pim_pz'],
                       m=PIMASS)

    v['slow'] = V4Array(r['slow_px'],
                        r['slow_py'],
                        r['slow_pz'],
                        m=PIMASS)
    
    return v    


def pr_dm(m1, m2, p1, p2):
    y = m1 - m2
    mag_p1 = p1.v.mag()
    mag_p2 = p2.v.mag()
    
    x = (mag_p1 - mag_p2) / (mag_p1 + mag_p2)
    
    return x, y


def pr_kpi_v_dm_pik_kpi(r, v):
    return pr_dm(r['m_pik'], r['m_kpi'], v['kp'], v['pim'])

def pr_pik_v_dm_pik_kpi(r, v):
    return pr_dm(r['m_pik'], r['m_kpi'], v['km'], v['pip'])
            

def p_kp(r, v):
    return v['kp'].v.mag()

def p_km(r, v):
    return v['km'].v.mag()

def p_pip(r, v):
    return v['pip'].v.mag()

def p_pim(r, v):
    return v['pim'].v.mag()

def p_slow(r, v):
    return v['slow'].v.mag()

def dm_pik_kpi(r, v):
    return r['m_pik'] - r['m_kpi']

def rs_dvf_m_pull(r, v):
    return (r['m_rs'] - paramdef['D0_MASS']) / r['rs_dvf_m_unc']

def ws_dvf_m_pull(r, v):
    return (r['m_ws'] - paramdef['D0_MASS']) / r['ws_dvf_m_unc']

column_calc = ColumnCalculator(
    four_vectors, ("kp_px kp_py kp_pz km_px km_py km_pz "
                   + "pip_px pip_py pip_pz pim_px pim_py pim_pz").split())

column_calc.add('pr_dm', pr_dm)
column_calc.add('pr_kpi_v_dm_pik_kpi', pr_kpi_v_dm_pik_kpi)
column_calc.add('pr_pik_v_dm_pik_kpi', pr_pik_v_dm_pik_kpi)
column_calc.add('p_kp', p_kp)
column_calc.add('p_pip', p_pip)
column_calc.add('p_km', p_km)
column_calc.add('p_pim', p_pim)
column_calc.add('p_slow', p_slow)
column_calc.add('dm_pik_kpi', dm_pik_kpi)
column_calc.add('rs_dvf_m_pull', rs_dvf_m_pull)
column_calc.add('ws_dvf_m_pull', ws_dvf_m_pull)
